// --*- C++ -*------x-----------------------------------------------------------
//
// Author:          Silvio Tosatto, Eckart Bindewald, Layla Hirsh
//
// Project name:    BiocompUP Align
//
// Date:            08/2005
//
// Description:     Compute Secondary Structure Element Alignment (SSEA)
//                  according to McGuffin & Jones.
//
// -----------------x-----------------------------------------------------------

#include <AGPFunction.h>
#include <Alignment.h>
#include <Debug.h>
#include <FSAlign.h>
#include <GetArg.h>
#include <NWAlign.h>
#include <ScoringS2S.h>
#include <SequenceData.h>
#include <SubMatrix.h>
#include <Substitution.h>
#include <SWAlign.h>
#include <Vec.h>
#include <algorithm>
#include <iostream>
#include <iomanip>
#include <math.h>
#include <queue>
#include <string>


using namespace Biopool;


void sAddLine()
{
	cout << "-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x" << endl;
}


void
helpOutput(ostream &os)
{
	os << "alignsec (version 2.1): Secondary structure alignment. Reads in 2"
	   << "(or more) secondary structures in FASTA format.\n"
	   << "The allowed character set is defined in the subsition matrix parameter"
	   << "file (default: alignsec2.prm).\n"
	   << "Usage: alignsec -a file1 -b file2\n"
	   << "Other options:\n"
	   << " -h                  : this info\n"
	   << " --matrix <filename> : specify substitution matrix (default: alignsec.prm)\n"
	   << " --sw                : if set, use local instead of global alignment\n"
	   << " --fs                : if set, use free-shift instead of global alignment\n"
	   << " --mult              : multiple alignment (database mode)\n"
	   << " --verbose 0|1|2     : set verbose level\n\n";
}


/// Get subset of set; ignore invalid indices of m.
template <class T>
inline vector<T>
getValidSubset(const vector<T> &v, const vector<int> &m)
{
	vector<T> result;

	for (unsigned int i = 0; i < m.size(); ++i)
		if ((m[i] >= 0) && (m[i] < static_cast<int>(v.size())))
			result.push_back(v[m[i]]);

	return result;
}


/**
 * Translate C, L and - into dot.
 */

/// Apply translation rule to string.
string
secstructConform(const string &sOrig, const string &rule)
{
	ERROR_IF(rule.size() != 2,
		"Translation rule must consist of 2 letters (fromChar+toChar).", exception);
	string s = sOrig;
	s = translate(s, rule[0], rule[1]);
	return s;
}


/// Apply set of translation rules to string. Each "rule" is a string with 2
/// characters, indicating the original charactor to be translated into the
/// second charactor of the rule example: .C is a rule to translate all dots
/// into the charactor C.
string
secstructConform(const string &sOrig, const vector<string> &rules)
{
	string s = sOrig;

	for (unsigned int i = 0; i < rules.size(); ++i)
		s = secstructConform(s, rules[i]);

	return s;
}


string
readSingleFasta(istream &is)
{
	string result;
	unsigned int remarks = 0;

	while (is)
	{
		string line = readLine(is);
		if (line.size() > 0)
			if (line[0] != '>')
				result = result + line;
			else
			{
				++remarks;
				if (remarks > 1)
					return result; // read only one remark line
			}
	}

	return result;
}


string
readFasta(istream &is, string &name)
{
	string result;
	unsigned int remarks = 0;

	while (is)
	{
		eatWhite(is);
		char c;
		is.get(c);
		if (!is) // get rid of EOF bug
			break;

		if (c == '>')
		{
			remarks++;
			if (remarks > 1)
			{
				is.putback(c); // put back char for next visit and exit
				return result;
			}
			else
				name = readLine(is);
		}
		else
		{
			string line = readLine(is);
			result = result + c + line; // c was taken in previous is.get()
		}
	}

	return result;
}



void
getSecondaryStructureBlockLengths(const string &sOrig, string &blocks,
	vector<unsigned int> &lengths, char coilChar)
{
	// Filter helices < 5 and strands < 3
	#ifdef DEBUG_VERBOSE
	cout << "Input sequence: " << sOrig << endl;
	#endif
	string s = sOrig;
	unsigned int count = 1;
	blocks = s[0];
	lengths.clear();
	for (unsigned int i = 1; i < s.size(); ++i)
	{
		if (s[i] == s[i-1])
			count++;
		else
		{
			// End of streak, check if too short
			ERROR_IF(count == 0,
				"Unexpected zero length in line 1631!", exception);
			lengths.push_back(count);
			blocks = blocks + s[i];
			count  = 1;
		}

	}
	lengths.push_back(count);
	#ifdef DEBUG_VERBOSE
	cout << "Blocks: " << blocks << "\n" << lengths << endl;
	#endif
	ERROR_IF(blocks.size() != lengths.size(),
		"Number of lengths was not equal number of blocks!", exception);
}


/// Bring word to size n by surrounding it with char c.
string
embedWord(const string &s, unsigned int n, char c, unsigned int left)
{
	ERROR_IF(n < s.size(),
		"n must be greater than word size!", exception);
	string result1 = s;
	if (n == s.size())
		return result1;
	string result;
	unsigned int diff = n - s.size();
	unsigned int right = diff - left;
	if (left > 0)
		result = result + string(left,c);
	result = result + s;
	if (right > 0)
		result = result + string(right,c);
	ERROR_IF(result.size() != n,
		"Internal error in line 1679!", exception);
	return result;
}


/// Return score of superposed strings according to matrix.
double
scoreMatch(const string &s1, const string &s2, const string &alphabet,
	const vector< vector<double> > &matrix)
{
	ERROR_IF(s1.size() != s2.size(),
		"Match strings must have equal size!", exception);
	#ifdef DEBUG_VERBOSE
	cout << "Starting scoreMatch!" << endl;
	#endif
	double result = 0;
	for (unsigned int i = 0; i < s1.size(); ++i)
	{
		unsigned int ui1 = alphabet.find(s1[i]);
		unsigned int ui2 = alphabet.find(s2[i]);
		if ((ui1 < matrix.size()) && (ui2 < matrix[ui1].size()))
			result += matrix[ui1][ui2];
	}
	#ifdef DEBUG_VERBOSE
	cout << "Ending scoreMatch!" << endl;
	#endif
	return result;
}


string
alignWord(const string &s1, const string &s2, const string &alphabet,
	const vector< vector<double> > &matrix)
{
	ERROR_IF(s1.size() > s2.size(),
		"Second string may not be shorter than first string!", exception);
	#ifdef DEBUG_VERBOSE
	cout << "Starting alignWord!\n"
	     << s1 << "\n"
	     << s2 << endl;
	#endif
	string bestString = s1;
	if (s1.size() == s2.size())
		return bestString;
	double bestScore = -9999.9;
	unsigned int diff = s2.size() - s1.size();
	for (unsigned int i = 0; i < diff; ++i)
	{
		string s = embedWord(s1, s2.size(), '-', i);
		double score = scoreMatch(s, s2, alphabet, matrix);
		if (score > bestScore)
		{
			bestScore = score;
			bestString = s;
		}
	}
	#ifdef DEBUG_VERBOSE
	cout << "Ending alignWord!\n"
	     << bestString << "\n"
	     << s2 << endl;
	#endif
	return bestString;
}


void
computeMatchStart(const vector<int> &matchInfo, const vector<unsigned int> &blockLenghts,
	unsigned int &counter, unsigned int &icounter)
{
	counter = 0;
	icounter = 0;
	unsigned int id = 0; // block id of first aligned block

	// Move to first position which is not "-"
	for (unsigned int i = 0; i < matchInfo.size(); ++i)
		if (matchInfo[i] != Align::INVALID_POS)
		{
			id = i;
			break;
		}
	for (unsigned int i = 0; i < static_cast<unsigned int>(matchInfo[id]); ++i)
	{
		ASSERT(i < blockLenghts.size(), exception);
		counter += blockLenghts[i];
	}

	icounter = matchInfo[id];
}


/// A low resolution (secondary structure) alignment is "extended" to a residue
/// resolution alignment. Very crude method! Only makes sense if no sequence
/// alignment is available at all.
///  @param s1  original sequence A
///  @param s2  original seqeuence B
///  @param sMatch1 result string of block alignment A
///  @param sMatch1 result string of block alignment B
///  @param l1 lengths of secondary structure elements of A
///  @param l2 lengths of secondary structure elements of B
///  @param accAlphabe alphabet of sequences
///  @param accMatrix substitution matrix for sequence alignments
vector<string>
computeExtendedMatch2(const string &s1, const string &s2,
                      const string &sMatch1, const string &sMatch2,
                      const vector<unsigned int> &l1, const vector<unsigned int> &l2,
                      const vector< vector<int> > &matchInfo,
                      const string &accAlphabet, const vector< vector<double> > &accMatrix,
                      unsigned int &startRes1, unsigned int &startRes2,
                      unsigned int &endRes1, unsigned int &endRes2)
{
	#ifdef DEBUG_VERBOSE
	cout << "Starting computeExtendedMatch2!\n"
	     << sMatch1 << "\n"
	     << sMatch2 << "\n"
	     << s1 << "\n"
	     << s2 << "\n"
	     << l1 << "\n"
	     << l2 << endl;
	#endif
	ERROR_IF(matchInfo.size() != 2,
		"Internal error: two match info vectors expected!", exception);
	vector<unsigned int> l1Subs = getValidSubset(l1, matchInfo[0]);
	vector<unsigned int> l2Subs = getValidSubset(l2, matchInfo[1]);
	ERROR_IF((sMatch1.size() != sMatch2.size()),
		"Match strings have uneven size!", exception);
	vector<string> result(2);
	vector<string> words1(sMatch1.size());
	vector<string> words2(sMatch2.size());
	unsigned int counter1 = 0; // counts current residue in original sequence
	unsigned int counter2 = 0;
	unsigned int icounter1 = 0; // counts current block
	unsigned int icounter2 = 0;
	computeMatchStart(matchInfo[0], l1, counter1, icounter1);
	computeMatchStart(matchInfo[1], l2, counter2, icounter2);
	startRes1 = counter1;
	startRes2 = counter2;


	for (unsigned int i = 0; i < sMatch1.size(); ++i)
	{
		if (sMatch1[i] == '-')
		{
			// Number of gap chars equal number l2[i]
			ERROR_IF(icounter2 >= l2.size(),
				"Too large icounter1", exception);
			ERROR_IF(l2[icounter2] == 0,
				"Unexected zero in line 1791", exception);
			words1[i] = string(l2[icounter2], '-');

		}
		else
		{
			ERROR_IF(icounter1 >= l1.size(),
				"Too large icounter1", exception);
			ERROR_IF(l1[icounter1] == 0,
				"Unexected zero in line 1791", exception);
			if ((counter1 + l1[icounter1] > s1.size()))
				words1[i] = s1.substr(counter1);
			else
				words1[i] = s1.substr(counter1, l1[icounter1]);

		}
		if (sMatch2[i] == '-')
		{
			ERROR_IF(icounter1 >= l1.size(),
				"Too large icounter1", exception);
			ERROR_IF(l1[icounter1] == 0,
				"Unexected zero in line 1791", exception);
			// Number of gap chars equal number l2[i]
			words2[i] = string(l1[icounter1], '-');
		}
		else
		{
			ERROR_IF(icounter2 >= l2.size(),
				"Too large icounter1", exception);
			ERROR_IF(l2[icounter2] == 0,
				"Unexected zero in line 1791", exception);
			if ((counter2 + l2[icounter2]) > s2.size())
				words2[i] = s2.substr(counter2);
			else
				words2[i] = s2.substr(counter2, l2[icounter2]);
		}
		if (sMatch1[i] != '-')
		{
			counter1 += words1[i].size();
			++icounter1;
		}
		if (sMatch2[i] != '-')
		{
			counter2 += words2[i].size();
			++icounter2;
		}
		if (words1[i].size() < words2[i].size())
			words1[i] = alignWord(words1[i], words2[i], accAlphabet, accMatrix);
		else
			if (words2[i].size() < words1[i].size())
				words2[i] = alignWord(words2[i], words1[i], accAlphabet, accMatrix);
		ERROR_IF(words1[i].size() != words2[i].size(),
			"Aligned words have uneven size!", exception);
	}

	// Construct result
	ERROR_IF(words1.size() != words2.size(),
		"Number of aligned words is not equal!", exception);
	for (unsigned int i = 0; i < words1.size(); ++i)
	{
		result[0] = result[0] + words1[i];
		result[1] = result[1] + words2[i];
	}
	endRes1 = startRes1 + result[0].size() - findPositions(result[0], '-').size() - 1; // counter1
	endRes2 = startRes2 + result[1].size() - findPositions(result[1], '-').size() - 1; // counter1
	ERROR_IF(result[0].size() != result[1].size(),
		"Result string have unequal size!", exception);
	#ifdef DEBUG_VERBOSE
	cout << "Finished computeExtendedMatch!\n"
	     << result[0] << "\n"
	     << result[1] << endl;
	#endif

	return result;
}


/// Main program.
int
main(int argc, char **argv)
{
	int gapPenalty = 0;
	char coilChar ='C'; // too short helices or strands are replaced to that char
	unsigned int startRes1 = 0;
	unsigned int startRes2 = 0;
	unsigned int endRes1 = 0;
	unsigned int endRes2 = 0;
	unsigned int verboseLevel = 1;

	vector<string> match;

	// Translate C, L and - into "."
	string targetName;
	string refName;

	string dataPath = getenv("VICTOR_ROOT");
	if (dataPath.length() < 3)
		cout << "Warning: environment variable VICTOR_ROOT is not set." << endl;
	dataPath.append("data/alignsec.prm");

	string secSubFileName = dataPath;

	string secBlocksTarget;
	string secBlocksTemplate;
	string secMatrixAlphabet;
	vector<unsigned int> secBlocksTargetLengths, secBlocksTemplateLengths;
	vector< vector<double> > secMatrix; // used for reading Blosum type matrix
	Vec<string> translationStartRules;
	Vec<string> translationEndRules;

	getArg("a", targetName, argc, argv, targetName);
	getArg("b", refName, argc, argv, refName);
	bool helpMode = getArg("h", argc, argv);
	bool smithWatermanMode = getArg("-sw", argc, argv);
	bool freeShiftMode = getArg("-fs", argc, argv);
	bool databaseMode = getArg("-mult", argc, argv);
	getArg("-matrix", secSubFileName, argc, argv, dataPath);
	getArg("-verbose", verboseLevel, argc, argv, 1);

	if (helpMode || (argc == 1))
	{
		helpOutput(cout);
		exit(0);
	}

	string path = getenv("VICTOR_ROOT");
	if (path.length() < 3)
		cout << "Warning: environment variable VICTOR_ROOT is not set." << endl;
	string examplesPath = path + "examples/";

	targetName = examplesPath + targetName;
	ifstream targetFile(targetName.c_str());
	ERROR_IF(!targetFile,
		"Error opening file with sequence 1: (option -a)", exception);

	refName = examplesPath + refName;
	ifstream refFile(refName.c_str());
	ERROR_IF(!refFile,
		"Error opening file with sequence 2: (option -b)", exception);

	string secTarget = readSingleFasta(targetFile);
	targetFile.close();

	// Read secondary structure alignment matrix
	ifstream secSubFile(secSubFileName.c_str());
	ERROR_IF(!secSubFile,
		"Error opening secondary structure matrix file.", exception);
	if (verboseLevel > 1)
		cout << "Reading secondary structure matrix file " << secSubFileName << endl;
	SubMatrix secSub(secSubFile);
	secSubFile.close();

	// Read same information again in different data structure:
	ifstream secSubFile2(secSubFileName.c_str());
	ERROR_IF(!secSubFile2,
		"Error opening secondary structure matrix file.", exception);
	secSubFile2 >> secMatrixAlphabet;
	Substitution::pReadDoubleVector(secSubFile2, secMatrix);

	// Basically read same information again
	secSubFile2 >> translationStartRules >> translationEndRules;
	secSubFile2.close();

	Align *ali;

	while (refFile)
	{
		string refTitle = "xxx";
		string refSec = readFasta(refFile, refTitle);

		if (!databaseMode)
			cout << "Sequence a:\n" << secTarget << "\n"
			     << "Sequence b:\n" << refSec << endl;
		else
		{
			sAddLine();
			cout << "QUERY Sequence:\n" << secTarget << "\n"
			     << "SUBJECT Sequence: " << refTitle << "\n" << refSec << endl;
		}

		getSecondaryStructureBlockLengths(refSec, secBlocksTemplate,
			secBlocksTemplateLengths, coilChar);

		ERROR_IF(secSub.size() == 0,
			"Secondary structure similarity matrix is not properly defined!", exception);

		secTarget = secstructConform(secTarget, translationStartRules);
		refSec = secstructConform(refSec, translationStartRules);

		if (verboseLevel > 1)
			cout << "Internal representation of sequence a:\n" << secTarget << "\n"
			     << "Internal representation of sequence b:\n" << refSec << "\n"
			     << "using matrix: " << secSub << " gap penalty: " << gapPenalty << endl;

		ERROR_IF(secSub.size() == 0,
			"Secondary structure matrix undefined.", exception);

		getSecondaryStructureBlockLengths(secTarget, secBlocksTarget,
			secBlocksTargetLengths, coilChar);

		SequenceData ad(2, secBlocksTarget, secBlocksTemplate);
		AGPFunction agp(gapPenalty, gapPenalty);
		ScoringS2S sc(&secSub, &ad, 0, 1.00);

		if (!smithWatermanMode)
		{
			cout << "Computing global alignment using secondary structure types and lengths.\n";
			ali = new NWAlign(&ad, &agp, &sc,
				secBlocksTargetLengths, secBlocksTemplateLengths);
		}
		else
			if (!freeShiftMode)
			{
				cout << "Computing local alignment using secondary structure types and lengths." << endl;
				ali = new SWAlign(&ad, &agp, &sc,
					secBlocksTargetLengths, secBlocksTemplateLengths);
			}
			else
			{
				cout << "Computing glocal alignment using secondary structure types and lengths." << endl;
				ali = new FSAlign(&ad, &agp, &sc,
					secBlocksTargetLengths, secBlocksTemplateLengths);
			}

		double meanLength = 0.5 * (secTarget.size() + refSec.size());
		double sovResult = 100.0 * 0.5 * ali->getScore();
		if (meanLength > 0.0)
		{
			sovResult = sovResult / meanLength;
		}
		cout << "Normalized score: " << sovResult << endl;
		if (verboseLevel > 1)
		{
			cout << "Alignment (internal representation):" << endl;
			ali->doMatch(cout);
			ali->doMatchScore(cout);
		}
		match = ali->getMatch(); // get alignment
		vector<string> matchExtern = match;
		for (unsigned int i = 0; i < matchExtern.size(); ++i)
			matchExtern[i] = secstructConform(matchExtern[i], translationEndRules);

		cout << "Block alignment:" << endl;
		for (unsigned int i = 0; i < matchExtern.size(); ++i)
			cout << matchExtern[i] << endl;

		ERROR_IF(!match.size() >= 2,
			"Internal error in line 619!", exception);
		vector<string> seqMatch = computeExtendedMatch2(secTarget, refSec,
			match[0], match[1],
			secBlocksTargetLengths,
			secBlocksTemplateLengths,
			ali->getMatchSubset(),
			secMatrixAlphabet, secMatrix,
			startRes1, startRes2,
			endRes1, endRes2);
		if (verboseLevel > 1)
			cout << "Experimental residue alignment (internal representation):\n"
			     << seqMatch[0] << "\n" << seqMatch[1] << endl;

		seqMatch[0] = secstructConform(seqMatch[0], translationEndRules);
		seqMatch[1] = secstructConform(seqMatch[1], translationEndRules);
		cout << "Experimental residue alignment:\n"
		     << startRes1 + 1 << "\t" << seqMatch[0] << "\t" << endRes1 + 1 << "\n"
		     << startRes2 + 1 << "\t" << seqMatch[1] << "\t" << endRes2 + 1 << endl;

		if (!databaseMode)
		{
			refFile.close();
			return 0;
		}

		delete ali;
	} // else: continue with next entry

	refFile.close();
	return 0;
}
